import os
from langchain_community.chat_models import ChatZhipuAI 
from dotenv import load_dotenv, find_dotenv
from langchain_core.prompts.chat import (
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
    AIMessagePromptTemplate,
    HumanMessagePromptTemplate,
)

_ = load_dotenv(find_dotenv())
model = os.environ.get('model')
if model is None:
    raise ValueError("model is not set in the .env file")

llm = ChatZhipuAI(model=model,
                  temperature=0.9,                 
    )

system_template = "你是一名实时翻译，你能将用户输入的消息同步翻译成英文和日文。"
system_message_prompt = SystemMessagePromptTemplate.from_template(
    system_template
    )

example_human = HumanMessagePromptTemplate.from_template("你好")

ai_template = "中文：你好。\n英文：Hello.日文：こんにちは。\n"
ai_message_prompt = AIMessagePromptTemplate.from_template(ai_template)

human_template="{text}"
human_message_prompt = HumanMessagePromptTemplate.from_template(
    human_template
    )

chat_prompt = ChatPromptTemplate.from_messages([
    system_message_prompt, 
    example_human, 
    ai_message_prompt, 
    human_message_prompt])

prompt = chat_prompt.invoke("我是只快乐的小蜜蜂。")
result = llm.invoke(prompt)

print(result.content)